"""
Mesh visualization utilities
"""
import numpy as np
try:
    import matplotlib.pyplot as plt
    from mpl_toolkits.mplot3d import Axes3D
    MATPLOTLIB_AVAILABLE = True
except ImportError:
    MATPLOTLIB_AVAILABLE = False

try:
    import vtk
    from vtk.util.numpy_support import vtk_to_numpy
    VTK_AVAILABLE = True
except ImportError:
    VTK_AVAILABLE = False

class MeshVisualizer:
    """Class for visualizing meshes and quality metrics"""
    
    @staticmethod
    def plot_mesh(mesh_data, title="Mesh Visualization", show=True, save_path=None):
        """
        Plot the mesh using matplotlib
        
        Parameters:
        -----------
        mesh_data : dict
            Dictionary containing mesh data
        title : str
            Plot title
        show : bool
            Whether to show the plot
        save_path : str, optional
            Path to save the plot
        """
        if not MATPLOTLIB_AVAILABLE:
            print("Matplotlib is not available. Install it with 'pip install matplotlib'")
            return
        
        points = mesh_data['points']
        cells = mesh_data['cells']
        
        # Create figure
        fig = plt.figure(figsize=(10, 8))
        ax = fig.add_subplot(111, projection='3d')
        
        # Plot mesh edges
        for cell in cells:
            if len(cell) == 3:  # Triangle
                for i in range(3):
                    p1 = points[cell[i]]
                    p2 = points[cell[(i + 1) % 3]]
                    ax.plot([p1[0], p2[0]], [p1[1], p2[1]], [p1[2], p2[2]], 'k-', alpha=0.5)
        
        # Set labels and title
        ax.set_xlabel('X')
        ax.set_ylabel('Y')
        ax.set_zlabel('Z')
        ax.set_title(title)
        
        # Set equal aspect ratio
        max_range = np.max([
            np.ptp(points[:, 0]),
            np.ptp(points[:, 1]),
            np.ptp(points[:, 2])
        ])
        mid_x = np.mean([np.min(points[:, 0]), np.max(points[:, 0])])
        mid_y = np.mean([np.min(points[:, 1]), np.max(points[:, 1])])
        mid_z = np.mean([np.min(points[:, 2]), np.max(points[:, 2])])
        ax.set_xlim(mid_x - max_range/2, mid_x + max_range/2)
        ax.set_ylim(mid_y - max_range/2, mid_y + max_range/2)
        ax.set_zlim(mid_z - max_range/2, mid_z + max_range/2)
        
        # Save if requested
        if save_path:
            plt.savefig(save_path, dpi=300, bbox_inches='tight')
        
        # Show if requested
        if show:
            plt.show()
        else:
            plt.close()
    
    @staticmethod
    def plot_quality_histogram(quality_metrics, title="Mesh Quality Histogram", show=True, save_path=None):
        """
        Plot a histogram of mesh quality values
        
        Parameters:
        -----------
        quality_metrics : dict
            Dictionary containing quality metrics
        title : str
            Plot title
        show : bool
            Whether to show the plot
        save_path : str, optional
            Path to save the plot
        """
        if not MATPLOTLIB_AVAILABLE:
            print("Matplotlib is not available. Install it with 'pip install matplotlib'")
            return
        
        quality_values = quality_metrics['quality_values']
        
        # Create figure
        plt.figure(figsize=(10, 6))
        
        # Plot histogram
        plt.hist(quality_values, bins=50, alpha=0.7, color='blue')
        
        # Add vertical lines for statistics
        plt.axvline(quality_metrics['min_quality'], color='r', linestyle='--', label=f"Min: {quality_metrics['min_quality']:.2f}")
        plt.axvline(quality_metrics['avg_quality'], color='g', linestyle='-', label=f"Avg: {quality_metrics['avg_quality']:.2f}")
        plt.axvline(quality_metrics['max_quality'], color='orange', linestyle='--', label=f"Max: {quality_metrics['max_quality']:.2f}")
        
        # Set labels and title
        plt.xlabel('Quality Value (Aspect Ratio)')
        plt.ylabel('Frequency')
        plt.title(title)
        plt.legend()
        plt.grid(True, alpha=0.3)
        
        # Save if requested
        if save_path:
            plt.savefig(save_path, dpi=300, bbox_inches='tight')
        
        # Show if requested
        if show:
            plt.show()
        else:
            plt.close()
    
    @staticmethod
    def plot_quality_comparison(original_metrics, optimized_metrics, title="Quality Comparison", show=True, save_path=None):
        """
        Plot a comparison of original and optimized mesh quality
        
        Parameters:
        -----------
        original_metrics : dict
            Dictionary containing original quality metrics
        optimized_metrics : dict
            Dictionary containing optimized quality metrics
        title : str
            Plot title
        show : bool
            Whether to show the plot
        save_path : str, optional
            Path to save the plot
        """
        if not MATPLOTLIB_AVAILABLE:
            print("Matplotlib is not available. Install it with 'pip install matplotlib'")
            return
        
        # Create figure
        plt.figure(figsize=(12, 6))
        
        # Plot histograms
        plt.hist(original_metrics['quality_values'], bins=50, alpha=0.5, color='red', label='Original')
        plt.hist(optimized_metrics['quality_values'], bins=50, alpha=0.5, color='green', label='Optimized')
        
        # Set labels and title
        plt.xlabel('Quality Value (Aspect Ratio)')
        plt.ylabel('Frequency')
        plt.title(title)
        plt.legend()
        plt.grid(True, alpha=0.3)
        
        # Add text with statistics
        improvement = (
            (original_metrics['avg_quality'] - optimized_metrics['avg_quality']) / 
            original_metrics['avg_quality'] * 100
        ) if original_metrics['avg_quality'] > 0 else 0
        
        stats_text = (
            f"Original Avg: {original_metrics['avg_quality']:.2f}\n"
            f"Optimized Avg: {optimized_metrics['avg_quality']:.2f}\n"
            f"Improvement: {improvement:.2f}%"
        )
        
        plt.annotate(
            stats_text, 
            xy=(0.05, 0.95), 
            xycoords='axes fraction',
            bbox=dict(boxstyle="round,pad=0.3", fc="white", ec="gray", alpha=0.8)
        )
        
        # Save if requested
        if save_path:
            plt.savefig(save_path, dpi=300, bbox_inches='tight')
        
        # Show if requested
        if show:
            plt.show()
        else:
            plt.close()
    
    @staticmethod
    def visualize_quality_vtk(mesh_data, quality_values=None, title="Mesh Quality Visualization"):
        """
        Visualize mesh quality using VTK
        
        Parameters:
        -----------
        mesh_data : dict
            Dictionary containing mesh data
        quality_values : numpy.ndarray, optional
            Array of quality values for each element
        title : str
            Visualization title
        """
        if not VTK_AVAILABLE:
            print("VTK is not available. Install it with 'pip install vtk'")
            print("Falling back to matplotlib visualization...")
            
            # Fall back to matplotlib visualization
            if quality_values is not None:
                # Create a dictionary with quality metrics
                quality_metrics = {
                    'quality_values': quality_values,
                    'min_quality': np.min(quality_values) if len(quality_values) > 0 else 0,
                    'max_quality': np.max(quality_values) if len(quality_values) > 0 else 0,
                    'avg_quality': np.mean(quality_values) if len(quality_values) > 0 else 0,
                    'std_quality': np.std(quality_values) if len(quality_values) > 0 else 0
                }
                
                # Plot quality histogram
                MeshVisualizer.plot_quality_histogram(quality_metrics, title=title)
            
            # Plot mesh
            MeshVisualizer.plot_mesh(mesh_data, title=title)
            return
        
        # Get or create VTK mesh
        if 'vtk_mesh' in mesh_data:
            mesh = mesh_data['vtk_mesh']
        else:
            # Create a new VTK mesh
            points = vtk.vtkPoints()
            for point in mesh_data['points']:
                points.InsertNextPoint(point)
            
            cells = vtk.vtkCellArray()
            for cell in mesh_data['cells']:
                if len(cell) == 3:  # Triangle
                    triangle = vtk.vtkTriangle()
                    triangle.GetPointIds().SetId(0, cell[0])
                    triangle.GetPointIds().SetId(1, cell[1])
                    triangle.GetPointIds().SetId(2, cell[2])
                    cells.InsertNextCell(triangle)
            
            mesh = vtk.vtkPolyData()
            mesh.SetPoints(points)
            mesh.SetPolys(cells)
        
        # Add quality values as cell data
        if quality_values is not None:
            quality_array = vtk.vtkDoubleArray()
            quality_array.SetName("Quality")
            for q in quality_values:
                quality_array.InsertNextValue(q)
            mesh.GetCellData().SetScalars(quality_array)
        
        # Create mapper and actor
        mapper = vtk.vtkPolyDataMapper()
        mapper.SetInputData(mesh)
        if quality_values is not None:
            mapper.SetScalarModeToUseCellData()
            mapper.SetScalarRange(np.min(quality_values), np.max(quality_values))
        
        actor = vtk.vtkActor()
        actor.SetMapper(mapper)
        
        # Create renderer and window
        renderer = vtk.vtkRenderer()
        renderer.AddActor(actor)
        renderer.SetBackground(0.1, 0.2, 0.4)
        
        render_window = vtk.vtkRenderWindow()
        render_window.AddRenderer(renderer)
        render_window.SetSize(800, 600)
        render_window.SetWindowName(title)
        
        # Create interactor
        interactor = vtk.vtkRenderWindowInteractor()
        interactor.SetRenderWindow(render_window)
        
        # Add color bar if quality values are provided
        if quality_values is not None:
            scalar_bar = vtk.vtkScalarBarActor()
            scalar_bar.SetLookupTable(mapper.GetLookupTable())
            scalar_bar.SetTitle("Quality")
            scalar_bar.SetNumberOfLabels(5)
            
            renderer.AddActor2D(scalar_bar)
        
        # Start interaction
        render_window.Render()
        interactor.Initialize()
        interactor.Start() 